#pragma once
#include <iostream>

namespace MYSPACE
{
  enum Color
  {
    RED = 0,
    BLACK
  };
  template <class T>
  class TreeNode
  {
  public:
    TreeNode(const T &data) : _parent(nullptr),
                              _left(nullptr),
                              _right(nullptr)
    {
      _data = data;
    }

  public:
    T _data;
    TreeNode *_parent;
    TreeNode *_left;
    TreeNode *_right;
    Color _color;
  };

  template <class T, class poi, class ref>
  class Iterator
  {
  public:
    typedef TreeNode<T> Node;
    typedef Iterator<T, poi, ref> Self;
    Iterator(Node *node) : _iterator(node) {}
    ref operator*()
    {
      return _iterator->_data;
    }
    poi operator->()
    {
      return &(_iterator->_data);
    }
    Self operator++()
    {
      Node *node = _iterator;
      //
      if (node->_right)
      {
        Node *left = node->_right;
        while (left->_left)
        {
          left = left->_left;
        }
        _iterator = left;
      }
      else
      {
        Node *parent = _iterator->_parent;
        while (node)
        {
          if (parent == nullptr)
          {
            _iterator = nullptr;
            break;
          }
          else if (parent->_left == node)
          {
            _iterator = parent;
            break;
          }
          else
          {
            node = parent;
            parent = parent->_parent;
          }
        }
      }
      return *this;
    }

    Self operator++(int)
    {
      Self copy = *this;
      ++(*this);
      return copy;
    }
    Self operator--()
    {
      Node *node = _iterator;
      //
      if (node->_left)
      {
        Node *right = node->_left;
        while (right->_right)
        {
          right = right->_right;
        }
        _iterator = right;
      }
      Node *parent = _iterator->_parent;
      while (node)
      {
        if (parent == nullptr)
        {
          _iterator = nullptr;
          break;
        }
        else if (parent->_right == node)
        {
          _iterator = parent;
          break;
        }
        else
        {
          node = parent;
          parent = parent->_parent;
        }
      }
      return *this;
    }
    Self operator--(int)
    {
      Self copy = *this;
      --(*this);
      return copy;
    }
    bool operator==(const Self &s)
    {
      return s._iterator == _iterator;
    }
    bool operator!=(const Self &s)
    {
      return s._iterator != _iterator;
    }

  private:
    Node *_iterator;
  };

  template <class K, class T, class KOfVCompare>
  class RBTree
  {
    typedef TreeNode<T> Node;

  public:
    typedef Iterator<T, T *, T &> iterator;
    typedef Iterator<const T, const T *, const T &> const_iterator;
    RBTree() : _root(nullptr) {}

  public:
    void print()
    {
      _print(_root);
    }
    bool IsRBTree()
    {
      int Bcount = 0, anycount = 0;
      Node *node = _root;
      while (node)
      {
        if (node->_color == BLACK)
        {
          Bcount++;
        }
        node = node->_left;
      }
      if (_IsRBTree(_root, Bcount, anycount))
      {
        return true;
      }
      else
      {
        return false;
      }
    }

  public:
    iterator begin()
    {
      Node *left = _root;
      while (left && left->_left)
      {
        left = left->_left;
      }
      return iterator(left);
    }
    iterator end()
    {
      return iterator(nullptr);
    }

    std::pair<iterator, bool> insert(const T &data)
    {
      if (_root == nullptr)
      {
        _root = new Node(data);
        _root->_color = BLACK;
        return std::make_pair(iterator(nullptr), false);
      }

      Node *parent = nullptr;
      Node *now_node = _root;
      while (now_node)
      {
        parent = now_node;
        if (_KOfV(now_node->_data) > _KOfV(data))
        // if(now_node->_data.first > data.first)
        {
          now_node = now_node->_left;
        }
        else if (_KOfV(now_node->_data) < _KOfV(data))
        // else if((now_node->_data.first) < (data.first))
        {
          now_node = now_node->_right;
        }
        else
        {
          return std::make_pair(iterator(now_node), false);
        }
      }
      Node *newnode = new Node(data);
      if (_KOfV(parent->_data) > _KOfV(data))
      // if((parent->_data.first) > (data.first))
      {
        parent->_left = newnode;
      }
      else if (_KOfV(parent->_data) < _KOfV(data))
      // else if((parent->_data.first) < (data.first))
      {
        parent->_right = newnode;
      }
      else
      {
        std::cerr << "error" << std::endl;
        exit(1);
      }
      newnode->_color = RED;
      newnode->_parent = parent;
      Node *copy = newnode;
      // 开始调整
      // 如果父亲颜色为黑 || parent== nullptr 则不管
      if (parent && parent->_color == RED)
      {
        Adjust(newnode);
      }
      return std::make_pair(copy, true);
    }
    iterator find(K key)
    {
      Node *node = _root;
      while (node)
      {
        if (_KOfV(node->_data) == key)
        {
          return iterator(node);
        }
        else if (_KOfV(node->_data) > key)
        {
          node = node->_left;
        }
        else
        {
          node = node->_right;
        }
      }
      return iterator(nullptr);
    }

  private:
    void Adjust(Node *node)
    {
      Node *parent = node->_parent;
      Node *grandparent = parent->_parent;
      Node *uncle = nullptr;
      // 如果父亲颜色为黑 || parent== nullptr 则不管
      while (parent && parent->_color == RED)
      {
        if (parent == grandparent->_left)
        {
          uncle = grandparent->_right;
          if (grandparent->_right == nullptr || uncle->_color == BLACK)
          {
            if (parent->_right == node)
            {
              // 变色双旋
              // 左右双旋
              LevoRotate(node);
              DextroRotate(node);
              grandparent->_color = RED;
              node->_color = BLACK;
              if (grandparent == _root)
              {
                _root = node;
              }
              break;
            }
            else if (parent->_left == node)
            {
              // 变色单旋
              // 右旋
              DextroRotate(parent);
              // adjust _color
              parent->_color = BLACK;
              grandparent->_color = RED;
              if (grandparent == _root)
              {
                _root = parent;
              }
              break;
            }
          }
          else if (uncle->_color == RED)
          {
            parent->_color = BLACK;
            uncle->_color = BLACK;
            grandparent->_color = RED;
          }
        }
        else if (parent == grandparent->_right)
        {
          uncle = grandparent->_left;
          if (grandparent->_left == nullptr || uncle->_color == BLACK)
          {
            if (parent->_right == node)
            {
              // 变色单旋
              // 左旋
              LevoRotate(parent);
              parent->_color = BLACK;
              grandparent->_color = RED;
              if (grandparent == _root)
              {
                _root = parent;
              }
              break;
            }
            else if (parent->_left == node)
            {
              // 变色双旋
              // 右左双旋
              DextroRotate(node);
              LevoRotate(node);
              grandparent->_color = RED;
              node->_color = BLACK;
              if (grandparent == _root)
              {
                _root = node;
              }
              break;
            }
          }
          else if (uncle->_color == RED)
          {
            parent->_color = BLACK;
            uncle->_color = BLACK;
            grandparent->_color = RED;
          }
        }
        node = grandparent;
        parent = node->_parent;
        if (parent)
          grandparent = parent->_parent;
      }
      if (_root)
      {
        _root->_color = BLACK;
      }
    }
    void DextroRotate(Node *node)
    {
      Node *parent = node->_parent;
      Node *grandparent = parent->_parent;
      Node *Rnode = node->_right;
      // adjust _left || _right
      node->_right = parent;
      parent->_left = Rnode;
      // adjust grandparent
      if (grandparent && grandparent->_left == parent)
      {
        grandparent->_left = node;
      }
      else if (grandparent && grandparent->_right == parent)
      {
        grandparent->_right = node;
      }
      // adjust _parent
      node->_parent = grandparent;
      parent->_parent = node;
      if (Rnode)
        Rnode->_parent = parent;
    }
    void LevoRotate(Node *node)
    {
      Node *parent = node->_parent;
      Node *grandparent = parent->_parent;
      Node *Lnode = node->_left;
      // adjust _left || _right
      node->_left = parent;
      parent->_right = Lnode;
      // adjust grandparent
      if (grandparent && grandparent->_left == parent)
      {
        grandparent->_left = node;
      }
      else if (grandparent && grandparent->_right == parent)
      {
        grandparent->_right = node;
      }
      // adjust _parent
      node->_parent = grandparent;
      parent->_parent = node;
      if (Lnode)
        Lnode->_parent = parent;
    }

  private:
    bool _IsRBTree(Node *root, int Bcount, int anycount)
    {
      if (root == nullptr)
      {
        if (Bcount != anycount)
        {
          std::cout << "Bcout:" << Bcount << " anycount:" << anycount << std::endl;
          return false;
        }
        else
        {
          return true;
        }
      }
      if (root->_color == BLACK)
      {
        ++anycount;
      }
      return _IsRBTree(root->_left, Bcount, anycount) &&
             _IsRBTree(root->_right, Bcount, anycount);
    }
    void _print(Node *root)
    {
      if (root == nullptr)
      {
        return;
      }
      _print(root->_left);
      // std::cout << _KOfV(root->_data) << std::endl;
      std::cout << (root->_data.first) << std::endl;
      _print(root->_right);
    }

  private:
    KOfVCompare _KOfV;
    Node *_root;
  };

}
